# amplitude.sessions -- session handling for WSGI apps
# by Sean A. Harrison <sah@tyndale.com>
# ---------------------------------------------------------------------------
# $Id: sessions.py 94 2013-07-31 20:39:13Z sah.harrison@gmail.com $

""" Session and SessionStorage classes
SessionStorage  : Base class, stores sessions in a dict in memory. VERY fast, but non-persistent.
FileStorage     : Stores sessions in a filesystem directory. Fast, persistent. Needs periodic cleanup.
DatabaseStorage : Stores sessions in a table in database. Fast, persistent. Needs periodic cleanup.
CookieStorage   : Stores sessions in the client's cookie, uses symmetric encryption. A bit slow.
MemcacheStorage : Stores sessions in memcached. Very fast. Needs periodic cleanup.
---------------------------------------------------------------------------
Memory Usage:                                         Footprint:
> python                                  1832 K
>>> from amplitude import sessions        2292 K      460 K
---------------------------------------------------------------------------
"""

DEBUG = False

import json
import os
from random import random
from time import time, asctime
from base64 import urlsafe_b64encode, urlsafe_b64decode
from amplitude.dict import Dict

# ---------------------------------------------------------------------------
class Session(Dict):
    """session is a dict, session.id is the id used in session storage
        Session(storage, **args) : returns a new session that will use the storage instance
        session[key]             : return the value of key in session
        session.id               : the unique id of the session
        session.save()           : save the session in storage
        Session.get(storage, id) : get the indicated session from storage.
                                    if it doesn't exist in storage, returns a new session
    Example:
    >>> storage = SessionStorage()          # new container for sessions
    >>> s = Session.get(storage)            # new session
    >>> s['name'] = 'amplitude'             # set a value
    >>> s.save()                            # save the session in storage
    >>> s1 = Session.get(storage, s.id)     # lookup the session by its id; use this with cookies.
    >>> s1.id == s.id and s1['name'] == 'amplitude'
    True
    >>> storage.delete(s.id)                # delete an existing session from storage
    >>> s2 = Session.get(storage, s.id)     # now s doesn't exist in storage, so s2 is a new session.
    >>> s2.id != s.id and s2.id != None and s2 == {}
    True
    """

    def __init__(self, storage, **args):
        self.id = self.new_key()
        self.storage = storage
        self.update(args)

    def new_key(self):
        """create a new session key"""
        # keep it simple but still pretty secure and guaranteed unique
        return urlsafe_b64encode("%s%.17f%.17f" % (time(), random(), random())).strip('=')[12:]

    @classmethod
    def get(cls, storage, id=None, **args):
        """get the session from storage by id. If a wrong id or None is given, return a new session."""
        if id:
            try:
                return storage.get(id)
            except KeyError:
                return Session(storage)
        else:
            return Session(storage)

    def save(self, **args):
        self.saved_at = asctime()
        self.storage.save(self, **args)
        if DEBUG==True: print "session saved:", self.id, self
        
    def delete(self):
        self.storage.delete(self.id)

# ---------------------------------------------------------------------------
class SessionStorage(dict):
    """A storage container for sessions. This base class implements memory storage.
    Example:
    >>> st = SessionStorage()
    >>> s = Session(st, name=u'sah')
    >>> s
    {'name': u'sah'}
    >>> s.save()
    >>> s1 = Session.get(st, s.id)
    >>> s1
    {'name': u'sah'}
    >>> s2 = Session.get(st, 'nonexistentid')
    >>> s2
    {}
    >>> s2.id == 'nonexistentid'
    False
    >>> st.delete(s.id)
    >>> s3 = Session.get(st, s.id)
    >>> s3 == {} and s3.id != s.id and s3.id != None
    True
    """

    def init_session(self, session_id=None, session_class=Session, **args):
        """returns a Session object initialized with self as SessionStorage and **args as keys.
        If session_id is given and exists in storage, that session is returned.
        """
        s = Session.get(self, session_id)
        s.storage = self
        s.update(args)
        return s

    # derived storage classes can override get() and save()
    def get(self, sessionid):
        """get the session associated with the given id from storage"""
        return self[sessionid]

    def save(self, session):
        """save the given session in storage"""
        self[session.id] = session

    def delete(self, sessionid):
        """delete the session associated with the given id from storage"""
        self.pop(sessionid, None)


class FileStorage(SessionStorage):
    """File-based storage for sessions.
        init with name of directory (which must exist) where sessions will be stored.
    Example: (pretty much the same as the example above)
    >>> st = FileStorage()
    >>> s = Session(st, name=u'sah')
    >>> s
    {'name': u'sah'}
    >>> s.save()
    >>> s1 = Session.get(st, s.id)
    >>> s1
    {u'name': u'sah'}
    >>> s2 = Session.get(st, 'nonexistentid')
    >>> s2
    {}
    >>> s2.id == 'nonexistentid'
    False
    >>> st.delete(s.id)
    >>> s3 = Session.get(st, s.id)
    >>> s3 == {} and s3.id != s.id and s3.id != None
    True
    """
    def __init__(self, directory='cache/sessions', **args):
        self.directory = os.path.abspath(directory)
        if not os.path.exists(self.directory): os.makedirs(self.directory)
        if not os.path.isdir(self.directory): raise "File '%s' exists but is not a directory" % directory

    def get(self, sessionid=None):
        if os.path.exists(self.session_file(sessionid)):
            f = open(self.session_file(sessionid))
            sdata = json.loads(f.read())
            f.close
            s = Session(self)
            s.update(sdata)
            s.id = sessionid
            return s
        else:
            return Session(self)

    def save(self, session):
        # make sure the directory for this session exists.
        if not os.path.exists(self.session_dir(session.id)):
            os.makedirs(self.session_dir(session.id))
        # write the file, overwriting what is there.
        f = open(self.session_file(session.id), 'w')
        f.write(json.dumps(session))
        f.close

    def delete(self, sessionid):
        if os.path.exists(self.session_file(sessionid)) \
        and not os.path.isdir(self.session_file(sessionid)):
            os.remove(self.session_file(sessionid))

    def session_dir(self, sessionid):
        """returns an absolute path to the session file."""
        # go 2 levels deep into self.directory, with 1000 x 1000 session folders,
        # making 1,000,000 session folders max.
        return os.path.join(self.directory, sessionid[0:3], sessionid[3:6])

    def session_file(self, sessionid):
        return os.path.join(self.session_dir(sessionid), sessionid)


class DatabaseStorage(SessionStorage):
    """Database storage of sessions.
        init with a DB-API 2.0 db connection.

    Example Session(requires empty 'test' db on localhost)
    >>> from amplitude import db
    >>> dbconn = db.Database('dbname=test user=postgres password=pgsql').connection
    >>> c = dbconn.cursor()
    >>> c.execute("create table sessions(id varchar primary key, data text)")
    >>> st = DatabaseStorage(dbconn)
    >>> s = Session(st, name=u'sah')
    >>> s
    {'name': u'sah'}
    >>> s.save()
    >>> s1 = Session.get(st, s.id)
    >>> s1
    {u'name': u'sah'}
    >>> s2 = Session.get(st, 'nonexistentid')
    >>> s2
    {}
    >>> s2.id == 'nonexistentid'
    False
    >>> st.delete(s.id)
    >>> s3 = Session.get(st, s.id)
    >>> s3 == {} and s3.id != s.id and s3.id != None
    True
    >>> c.execute("drop table sessions")
    >>> dbconn.commit()
    """

    def __init__(self, dbconn, table='sessions', idfield='id', datafield='data', **args):
        self.dbconn = dbconn
        self.table = table
        self.idfield = idfield
        self.datafield = datafield

    def get(self, sessionid=None):
        c = self.dbconn.cursor()
        c.execute("select %s from %s where %s='%s' limit 1" % (self.datafield, self.table, self.idfield, sessionid))
        res = c.fetchone()
        c.close()
        if res:
            sdata = json.loads(res[0])
            s = Session(self)
            s.update(sdata)
            s.id = sessionid
            return s
        else:
            return Session(self)

    def save(self, session):
        sdata = json.dumps(session)
        c = self.dbconn.cursor()
        try:
            c.execute("delete from %s where %s='%s'" % (self.table, self.idfield, session.id))
            c.execute("insert into %s (%s, %s) values ('%s', '%s')" % (self.table, self.idfield, self.datafield, session.id, sdata))
            self.dbconn.commit()
        except:
            self.dbconn.rollback()
            raise
        c.close()

    def delete(self, sessionid):
        c = self.dbconn.cursor()
        try:
            c.execute("delete from %s where %s='%s'" % (self.table, self.idfield, sessionid))
            self.dbconn.commit()
        except:
            self.dbconn.rollback()
            raise
        c.close()


class CookieStorage(SessionStorage):
    """Cookie storage of sessions.
        init with an encryption key, and request and response objects.

    Example Session (simulating Request and Response objects)
    >>> from amplitude.http import Request, Response
    >>> st = CookieStorage('thiskeyisverybad,bettergetrandom')
    >>> s = Session(st, name=u'sah')
    >>> s
    {'name': u'sah'}
    >>> s.save()
    >>> s1 = Session.get(st, s.id)
    >>> s1
    {u'name': u'sah'}
    >>> s2 = Session.get(st, 'nonexistentid')
    >>> s2
    {}
    >>> s2.id == 'nonexistentid'
    False
    >>> st.delete(s.id)             # <-- This has no effect, there is no concept of deleting a cookie session.
    >>> s3 = Session.get(st, s.id)  # <-- The session data is actually in s.id, so the following tests are all false.
    >>> s3 == {} and s3.id != s.id and s3.id != None
    False
    """

    def __init__(self, key):
        from crypto.cipher.aes import AES
        self.key = key
        self.cipher = AES(self.key, keySize=len(key))

    def get(self, sessionid=None):
        # with a cookie session, the sessionid is actually the encrypted data for the session
        session = Session(self)
        if sessionid:
            try:
                sdata = self.__decrypt(sessionid)
            except:
                sdata = {}
            session.update(sdata)
        session.id = self.__encrypt(session)
        return session

    def save(self, session):
        session.id = self.__encrypt(session)

    def delete(self, sessionid, **args):
        pass

    def __encrypt(self, session):
        return urlsafe_b64encode(self.cipher.encrypt(json.dumps(session)))

    def __decrypt(self, sessionid):
        return json.loads(self.cipher.decrypt(urlsafe_b64decode(sessionid)))

    @classmethod
    def make_key(cls, keySize=32):
        base92 = " 0 1 2 3 4 5 6 7 8 9 a b c d e f g h i j k l m n o p q r s t u v w x y z A B C D E F G H I J K L M N O P Q R S T U V W X Y Z ` ~ ! @ # $ % ^ & * ( ) _ + - = [ { ] } \ | ; : , < . > / ? ".split()
        t = []
        for i in range(keySize):
            t.append(base92[int(random()*92)])
        return "".join(t)


class MemcacheStorage(SessionStorage):
    """Memcache storage of sessions.
        init with memcache server location (default 127.0.0.1:11211)

    Example Session
    >>> st = MemcacheStorage('127.0.0.1:11211')
    >>> s = Session(st, name=u'sah')
    >>> s
    {'name': u'sah'}
    >>> s.save()
    >>> s1 = Session.get(st, s.id)
    >>> s1
    {u'name': u'sah'}
    >>> s2 = Session.get(st, 'nonexistentid')
    >>> s2
    {}
    >>> s2.id == 'nonexistentid'
    False
    >>> st.delete(s.id)
    >>> s3 = Session.get(st, s.id)
    >>> s3 == {} and s3.id != s.id and s3.id != None
    True
    """

    def __init__(self, serverlocation='127.0.0.1:11211'):
        import memcache
        self.client = memcache.Client([serverlocation])

    def get(self, sessionid=None):
        session = Session(self)
        if sessionid:
            sdata = self.client.get(sessionid)
            if sdata:
                session.update(json.loads(sdata))
                session.id = sessionid
        else:
            self.save(session)
        return session

    def save(self, session):
        self.client.set(session.id, json.dumps(session))
        self[session.id] = {'saved': session.saved_at}

    def delete(self, sessionid):
        self.client.delete(sessionid)

# ---------------------------------------------------------------------------

def doctests():
    """ -- Additional doctests --
    """

if __name__ == "__main__":
    import doctest
    doctest.testmod()
